import pysam
import argparse
from tqdm import tqdm
from deepsignal3.utils.process_utils import get_refloc_of_methysite_in_motif
from deepsignal3.utils.process_utils import get_motif_seqs
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inbam','-i', type=str, default="/public/data/wangjianxin/nipeng/workspace/arab_col0_csu2_nextomics.ont_r10.4.1_simplex.pod5_pass.bam", required=False,
                        help="bam ready to be split")
    parser.add_argument('--out', '-o', type=str, default="/public/home/hpc234701005/data/R10.4/arab/train/arab_col0_train.bam", required=False,
                        help="where the splited bam to store")
    parser.add_argument('--pod5_dir', type=str, required=False,
                        help="id")
    parser.add_argument('--tsv', type=str, required=False,
                        help="id")
    parser.add_argument('--stat', action='store_true', default=False, help="stat mod")
    parser.add_argument(
        "--motifs",
        action="store",
        type=str,
        required=False,
        default="CG",
        help="motif seq to be extracted, default: CG. "
        "can be multi motifs splited by comma "
        "(no space allowed in the input str), "
        "or use IUPAC alphabet, "
        "the mod_loc of all motifs must be "
        "the same",
    )
    parser.add_argument(
        "--mod_loc",
        action="store",
        type=int,
        required=False,
        default=0,
        help="0-based location of the targeted base in the motif, default 0",
    )
    return parser.parse_args()

def get_q2tloc_from_cigar(r_cigar_tuple, strand, seq_len):
    """
    insertion: -1, deletion: -2, mismatch: -3
    :param r_cigar_tuple: pysam.alignmentSegment.cigartuples
    :param strand: 1/-1 for fwd/rev
    :param seq_len: read alignment length
    :return: query pos to ref pos
    """
    fill_invalid = -2
    # get each base calls genomic position
    q_to_r_poss = np.full(seq_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar_tuple if strand == 1 else r_cigar_tuple[::-1]
    for op, op_len in cigar_ops:
        if op == 1:
            # inserted bases into ref
            for q_pos in range(curr_q_pos, curr_q_pos + op_len):
                q_to_r_poss[q_pos] = -1
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                q_to_r_poss[curr_q_pos + op_offset] = curr_r_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    q_to_r_poss[curr_q_pos] = curr_r_pos
    if q_to_r_poss[-1] == fill_invalid:
        raise ValueError(
            (
                "Invalid cigar string encountered. Reference length: {}  Cigar "
                + "implied reference length: {}"
            ).format(seq_len, curr_r_pos)
        )
    return q_to_r_poss

def get_parent_id(bam_read):
    try:
        # if pi tag is present this is a child read
        return bam_read.get_tag("pi")
    except KeyError:
        # else this is the parent read so return query_name
        return bam_read.query_name

def deal_bam(args):
    input_bam = pysam.AlignmentFile(args.inbam, 'rb',check_sq=False)
    output_file=open(args.out, 'w')
    motif_seqs = get_motif_seqs(args.motifs, True)
    for bam_read in input_bam:
        if bam_read.is_supplementary or bam_read.is_secondary:
            continue
        reference_name = bam_read.reference_name
        if reference_name is None:
            continue
        read_name = bam_read.query_name#get_parent_id(bam_read)
        #ref_loc = bam_read.get_reference_positions(full_length=True)
        strand = "-" if bam_read.is_reverse else "+"
        seq = bam_read.get_forward_sequence()
        ref_start = bam_read.reference_start
        ref_end = bam_read.reference_end
        strand_code = 0 if bam_read.is_reverse else 1
        cigar_tuples = bam_read.cigartuples
        tsite_locs = get_refloc_of_methysite_in_motif(
                                    seq, set(motif_seqs), args.mod_loc)
        qalign_start = bam_read.query_alignment_start
        qalign_end = bam_read.query_alignment_end
        if bam_read.is_reverse:
            seq_start = len(seq) - qalign_end
            seq_end = len(seq) - qalign_start
        else:
            seq_start = qalign_start
            seq_end = qalign_end
        q_to_r_poss = get_q2tloc_from_cigar(
            cigar_tuples, strand_code, (seq_end - seq_start)
        )
        # ref_readlocs = dict()
        # ref_poss = []
        # pred_pos = []
        for loc_in_read in tsite_locs:
            if seq_start <= loc_in_read < seq_end:
                offset_idx = loc_in_read - seq_start
                if q_to_r_poss[offset_idx] != -1:
                    if strand == "-":
                        # pos = '.'#loc_in_read
                        ref_pos = ref_end - 1 - \
                            q_to_r_poss[offset_idx]
                    else:
                        # pos = loc_in_read
                        ref_pos = ref_start + \
                            q_to_r_poss[offset_idx]
                    # ref_readlocs[loc_in_read] = ref_pos
                    # ref_poss.append(ref_pos)
                    # pred_pos.append(loc_in_read)
                    output_file.write(read_name+'\t' + str(loc_in_read) + '\t'+ reference_name + '\t' + str(ref_pos) + '\n')
                else:
                    continue
                    # ref_pos = -1
                    # ref_readlocs[loc_in_read] = ref_pos
                    # ref_poss.append(ref_pos)
                    # pred_pos.append(loc_in_read)
            else:
                continue
                # ref_pos = -1
                # ref_readlocs[loc_in_read] = ref_pos
                # ref_poss.append(ref_pos)
                # pred_pos.append(loc_in_read)


    input_bam.close()
    output_file.close()
if __name__ == '__main__':
    args=parse_args()
    deal_bam(args)